import pandas as pd
from sklearn.tree import DecisionTreeClassifier
testfile = './tmp/test_cleaned.csv' # 已预处理和贴标签训练数据
trainfile_class = './tmp/train_class.csv' # 已预处理测试数据
train_class = pd.read_csv(trainfile_class)
test = pd.read_csv(testfile)

# 删除列
x_train = train_class.drop(['user_id','merchant_id','coupon_id','date_received','date'],axis=1)
x_test = test.drop(['user_id','merchant_id','coupon_id','date_received','date'],axis=1)

# 决策树建模
model_dt1 = DecisionTreeClassifier(max_leaf_nodes=16, random_state=123).fit(
    x_train.iloc[:, :-1], x_train.iloc[:, -1]
)

# 预测结果
pre_dt = model_dt1.predict(x_test.values)

# DataFrame存放决策预测结果
dt_class = test[['user_id', 'merchant_id', 'coupon_id']]
dt_class['class'] = pre_dt
dtfile_pre = './tmp/dt_class.csv' # 导出数据
dt_class.to_csv(dtfile_pre, index=False)

